/*
  Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Functions for reading GADGET HDF5 snapshot files.

#include "config.h"

#if defined (HAVE_LIBHDF5) && (HAVE_LIBHDF5_CPP) && (HAVE_H5CPP_H)

#include "H5Cpp.h"
#include "snapshot.h"
#include <string>
#include <numeric>
#include "boost/lexical_cast.hpp"
#include "preferences.h"
#include "constants.h"
#include "boost/shared_ptr.hpp"
#include "boost/lambda/lambda.hpp"
#include "boost/lambda/bind.hpp"
#include "mcrx-debug.h"
#include "mcrx-types.h"

using namespace std;
using namespace H5;
using namespace blitz;
using namespace boost;
using namespace boost::lambda;
using boost::shared_ptr;

template <typename T> class h5Typemapper;

template <> class h5Typemapper<double> {
public:
  static const PredType h5type;
};
const PredType h5Typemapper<double>::h5type(PredType::NATIVE_DOUBLE);
template <> class h5Typemapper<float> {
public:
  static const PredType h5type;
};
const PredType h5Typemapper<float>::h5type(PredType::NATIVE_FLOAT);
template <> class h5Typemapper<int> {
public:
  static const PredType h5type;
};
const PredType h5Typemapper<int>::h5type(PredType::NATIVE_INT);
template <> class h5Typemapper<unsigned int> {
public:
  static const PredType h5type;
};
const PredType h5Typemapper<unsigned int>::h5type(PredType::NATIVE_UINT);

template <typename T> 
void readValue(const Attribute& a, T& value)
{
  a.read(h5Typemapper<T>::h5type, &value);
}
 
class unexpectedRank {
public:
  int rank_;
  unexpectedRank(int r) : rank_(r) {};
};

class unexpectedWidth {
public:
  int width_;
  unexpectedWidth(int w) : width_(w) {};
};

/** Read a 1D dataset. */
template <typename T_field>
bool 
mcrx::Snapshot::read_dataset (const CommonFG& f, const string& datasetname,
			      vector<T_field>& field,
			      const vector<int>& read_order) const
{
  DataSet dataset;
  try {
    dataset = f.openDataSet(datasetname );
  }
  catch (DataSetIException&) {
    DEBUG(1, cout << "Dataset not found: " << datasetname << endl;);
    // if we couldn't read it, that just means there are no particles
    return false;
  }
  catch (GroupIException&) {
    DEBUG(1, cout << "Nonexisting group reading dataset: " << datasetname << endl;);
    // if we couldn't read it, that just means there are no particles
    return false;
  }
      
  DataSpace dataspace = dataset.getSpace();
  int rank = dataspace.getSimpleExtentNdims();
  if(rank>1) {
    // the field does not have the dimensionality we expected
    throw unexpectedRank(rank);
  }
  hsize_t extent;
  dataspace.getSimpleExtentDims( &extent, NULL);

  DEBUG(1, cout << "Reading dataset: " << datasetname << " size " << extent << endl;);

  field.resize(extent);
  if (!read_order.empty()) {
    // read data
    vector<T_field> temp(extent);
    dataset.read(&temp[0], h5Typemapper<T_field>::h5type,
		 dataspace, dataspace );

    // put in correct order
    for(int i=0;i<temp.size(); ++i) {
      assert (read_order [i] < field.size());
      field [read_order [i]] = temp[i];
    }
  }
  else {
    // no reonrdering, we can read directly into vector
    dataset.read(&field[0], h5Typemapper<T_field>::h5type,
		 dataspace, dataspace );
  }
  return true;
}

/** Read a 2D dataset. */
template<typename T, int N>
bool 
mcrx::Snapshot::read_dataset (const CommonFG& f, const string& datasetname,
			      vector<TinyVector<T, N> > & field,
			      const vector<int>& read_order) const
{
  DataSet dataset;
  try {
    dataset = f.openDataSet(datasetname );
  }
  catch (DataSetIException&) {
    DEBUG(1, cout << "Dataset not found: " << datasetname << endl;);
    // if we couldn't read it, that just means there are no particles
    return true;
  }
  catch (GroupIException&) {
    DEBUG(1, cout << "Nonexisting group reading dataset: " << datasetname << endl;);
    // if we couldn't read it, that just means there are no particles
    return true;
  }
  DataSpace dataspace = dataset.getSpace();
  int rank = dataspace.getSimpleExtentNdims();
  if(rank!=2) {
    // the field does not have the dimensionality we expected
    throw unexpectedRank(rank);
  }
  blitz::TinyVector<hsize_t, 2> extent;
  dataspace.getSimpleExtentDims( &extent[0], NULL);
  if(extent[1]!=N) {
    throw unexpectedWidth(extent[1]);
  }

  DEBUG(1, cout << "Reading dataset: " << datasetname << " size " << extent << endl;);

  field.resize(extent[0]);
  if (!read_order.empty()) {
    // read data. because tinyvectors are padded due to alignment
    // constraints, we read into an array and then copy them one at a time.

    //vector<TinyVector<T,N> > temp(extent[0]);
    array_2 temp(extent[0],N, blitz::contiguousArray);
    dataset.read(temp.dataFirst(), h5Typemapper<T>::h5type,
		 dataspace, dataspace );

    // put in correct order
    for(int i=0;i<field.size(); ++i) {
      assert (read_order [i] < field.size());
      field [read_order [i]] = temp(i,Range::all());
    }
  }
  else {
    // no reordering, we can read directly into vector
    dataset.read(&field[0], h5Typemapper<T>::h5type,
		 dataspace, dataspace );
  }

  return true;
}


inline double mcrx::Snapshot::a2time(double a, double OmegaM, double OmegaL, 
				     double UnitTime_in_s)
{
  // only works for flat
  assert(abs(OmegaM+OmegaL-1)<1e-6);
  const double Hubble_in_cgs            = 3.2407789e-18;   // h/sec
  const double Hubble = Hubble_in_cgs  * UnitTime_in_s;
  const double UnitTime_in_Gyr = UnitTime_in_s / (3600.*24.*365.*1e9);
  const double aml = pow(OmegaM/OmegaL, 1./3);
  const double time = (2.0/(Hubble* 3.* sqrt(1.-OmegaM))*
		       log(pow(a/aml,1.5) + sqrt(1+pow(a/aml,3.))))
    * UnitTime_in_Gyr;
  return time;
}
    

/** Loads a GADGET snapshot HDF5 file.  */
bool mcrx::Snapshot::loadgadget_hdf5(const string& snapfn)
{
  cout << "Loading Gadget HDF5 snapshot file: " << snapfn << endl;

  Exception::dontPrint();
  H5open();
  H5File file;
  try {
    file.openFile( snapfn, H5F_ACC_RDONLY );
  }
  catch (H5::FileIException&) {
    cerr << "Error: could not open file" << endl;
    throw;
  }

  const Group header = file.openGroup("Header");

  double OmegaM, OmegaL, h0;
  bool comoving_units=false;

  // if comovingIntegration is on, time is actually a. we can't handle that
  if(prefs_.getValue("ComovingIntegrationOn",int())) {
    comoving_units=true;
    readValue(header.openAttribute("Time"), a_);
    cout << "  Comoving snapshot, expansion factor is " << a_ << endl;
    cout << "  converting expansion factors to time" << endl;
    readValue(header.openAttribute("HubbleParam"), h0);
    readValue(header.openAttribute("Omega0"), OmegaM);
    readValue(header.openAttribute("OmegaLambda"), OmegaL);
  }
  else {
    a_=1;
    readValue(header.openAttribute("Time"), time_);
  }

  // Get units from Gadget parameter file (NOTE: UNITS ARRIVE AS cgs units)
  // length in kpc
  lengthunit_ = a_*(prefs_.getValue("UnitLength_in_cm",double()))/
    (constants::kpc*1e2);
  // mass in Msun
  massunit_ = (prefs_.getValue("UnitMass_in_g",double()))/
    (constants::Msun*1e3);
  // time in yr
  timeunit_ = prefs_.getValue("UnitLength_in_cm",double())/(prefs_.getValue("UnitVelocity_in_cm_per_s",double())*constants::yr);

  const double UnitTime_in_s = prefs_.getValue("UnitLength_in_cm",double())/
    prefs_.getValue("UnitVelocity_in_cm_per_s",double());

  if(user_prefs_.defined("h-inverse_units") && 
     user_prefs_.getValue("h-inverse_units", bool())) {
    // all numbers should be multiplied by h^-1. This means that we
    // can divide the units by h.
    readValue(header.openAttribute("HubbleParam"), h0);
    cout << "  Using h^-1 units with h=" << h0 << endl;
    lengthunit_ /= h0;
    massunit_ /= h0;
    timeunit_ /= h0;
  }

  G_=0; // Note: _G will be incorrect until makephysicalunits() is run, we can figure out a _G here and we should. don't bother, we don't use G.
  units_["length"] = boost::lexical_cast<string>(lengthunit_)+"*kpc";
  units_["mass"] = boost::lexical_cast<string>(massunit_)+"*Msun";
  units_["time"] = boost::lexical_cast<string>(timeunit_)+"*yr";
  units_["temperature"] = "K";

  if(comoving_units)
    time_=a2time(a_, OmegaM, OmegaL, UnitTime_in_s);

  cout << "  Snapshot time is " << time_ << " (" 
       << time_*timeunit_ << " yr)" <<endl;

  // syntactic convenience. note that we can't use the n(species)
  // function yet, because it relies on the particle data being loaded.
  vector<int> npart(n_species);
  // NOT use numpart_total, thatis across all snapshots
  readValue(header.openAttribute("NumPart_ThisFile"), npart[0]);
  //const int n_tot = accumulate(npart.begin(), npart.end(), 0);
  cout << "  Number of particles: ";
  for(int i=0; i<n_species; ++i)
    cout << npart[i] <<' ';
  cout << '\n';
  // for convenience, set up an array of the hdf5 groups also, if we
  // have selected to not read any of halo/disk/bulge then we can
  // easily do that here
  vector<boost::shared_ptr<Group> > groups;
  for(int s=0; s<n_species; ++s) {
    if ((npart[s]>0) && 
	(!user_prefs_.defined("load_"+species_names_[s]) ||
	 user_prefs_.getValue("load_"+species_names_[s], bool()) ))
      groups.push_back(boost::shared_ptr<Group>(new Group(file.openGroup("PartType"+lexical_cast<string>(s)))));
    else {
      groups.push_back(boost::shared_ptr<Group>(new Group()));
      npart[s]=0;
    }
  }


  // First set up ordering maps and shuffle container. See loadgadget for info.
  vector<map<unsigned int, int> > shuffle(n_species);
  vector<vector<int> > read_order(n_species);
  for(int s=0; s < n_species; ++s) {
    // read data for shuffle map
    vector<unsigned int> id;
    read_dataset (*groups[s], "ParticleIDs", id);
    assert(id.size() == npart[s]);

    // Set up shuffle map for species s 
    // we can NOT just put the id in the vector here!
    for(int i=0;i<npart[s]; ++i) {
      shuffle[s][id[i]] = i; // ID->file pos
    }

    // Check that all ID numbers are unique (at least for a given species)
    if(shuffle[s].size() < npart[s]) {
      cerr << "FATAL: Duplicate ID numbers in species " << s << ": " << shuffle[s].size() << '\t' << npart[s] << endl;
      throw particle_set_generic::illegal_id();
    }

    setup_ordering(read_order[s], *sets()[s], shuffle[s], npart[s]);

    // NOW we can put the id numbers into the grid
    sets()[s]->id_.resize(npart[s]);
    for(int i=0;i<npart[s]; ++i) {
      sets()[s]->id_[read_order[s][i]] = id[i];
    }
  }

  //test ordering
  {
    vector<unsigned int> id;
    read_dataset (*groups[0], "ParticleIDs", id, read_order[0]);
    for(int i=0; i<npart[0]; ++i) assert(id[i]==sets()[0]->id_[i]);
  }
    

  /// now it's easy, crank away at reading
  /// \todo this should probably be delegated to the particle set itself...
  
  // Read basic properties: positions & Velocities
  vector<double> masses(n_species);
  readValue(header.openAttribute("MassTable"), masses[0]);
  for(int s=0; s < n_species; ++s) {
    read_dataset (*groups[s], "Coordinates", sets()[s]->pos_, read_order[s]);
    read_dataset (*groups[s], "Velocities", sets()[s]->vel_, read_order[s]);
    if(masses[s]==0)
      read_dataset (*groups[s], "Masses", sets()[s]->m_, read_order[s]);
    else {
      sets()[s]->m_.assign(npart[s], masses[s]);
      DEBUG(1, cerr << "Constant mass for "<<npart[s]<< " particles"<<endl;);
    }
  }
    
  // read gas quantities
  read_dataset (*groups[gas], "SmoothingLength", 
		sets()[gas]->h_, read_order[gas]);
  read_dataset (*groups[gas], "StarFormationRate", 
		gas_particles.sfr_, read_order[gas]);
  vector<float> u_ngas;
  read_dataset (*groups[gas], "InternalEnergy", 
		u_ngas, read_order[gas]);
  vector<float> ne;
  read_dataset (*groups[gas], "ElectronAbundance", 
		ne, read_order[gas]);
  vector<float> tp;
  read_dataset (*groups[gas], "EnergyReservoirForFeeback", 
		tp, read_order[gas]);
  if(tp.empty()) {
    cerr << "EnergyReservoirForFeedback not found -- effective pressure will be the hydrodynamic pressure\n";
    tp.assign(ne.size(), 0.0);
  }
    
  // Calculate temp and teff
  DEBUG(1,cout << "Calculating effective temperatures" << endl;);
  calculate_temperatures(u_ngas,ne,tp);

  // we init the disk and bulge particles here so they can be
  // overwritten below if their metallicities and formation times
  // exist in the file
  init_disk_bulge_halo();

  // try to read metallicity as a scalar field
  try {
    read_dataset (*groups[gas], "Metallicity", 
		  gas_particles.z_, read_order[gas]);
    read_dataset (*groups[disk], "Metallicity", 
		  disk_particles.z_, read_order[disk]);
    read_dataset (*groups[bulge], "Metallicity", 
		  bulge_particles.z_, read_order[bulge]);
    read_dataset (*groups[nstar], "Metallicity", 
		  nstar_particles.z_, read_order[nstar]);
  }
  catch (unexpectedRank& r) {
    //  this means it's a vector field, ie Scannapieco metals
    vector<TinyVector<float,12> > metals;
    read_dataset (*groups[gas], "Metallicity", 
		  metals, read_order[gas]);
    decode_scannapieco(metals, gas_particles);

    read_dataset (*groups[disk], "Metallicity", 
		  metals, read_order[disk]);
    decode_scannapieco(metals, disk_particles);

    read_dataset (*groups[bulge], "Metallicity", 
		  metals, read_order[bulge]);
    decode_scannapieco(metals, bulge_particles);

    read_dataset (*groups[nstar], "Metallicity", 
		  metals, read_order[nstar]);
    decode_scannapieco(metals, nstar_particles);
  }
  // back-convert SFR to gadget units so that it can get converted in
  // the general setup routine (this is retarded)
  transform (gas_particles.sfr_.begin(),gas_particles.sfr_.end(),
	     gas_particles.sfr_.begin(),
	     _1 * (timeunit_/massunit_));

  // formation times
  read_dataset (*groups[disk], "StellarFormationTime", 
		disk_particles.tf_, read_order[disk]);
  read_dataset (*groups[bulge], "StellarFormationTime", 
		bulge_particles.tf_, read_order[bulge]);
  read_dataset (*groups[nstar], "StellarFormationTime", 
		nstar_particles.tf_, read_order[nstar]);

  if(comoving_units) {
    // convert expansion factors to times. these times still have h^-1
    // in them so they get converted later
    transform(disk_particles.tf_.begin(),
	      disk_particles.tf_.end(),
	      disk_particles.tf_.begin(),
	      bind(&a2time, _1, OmegaM, OmegaL, UnitTime_in_s));
    transform(bulge_particles.tf_.begin(),
	      bulge_particles.tf_.end(),
	      bulge_particles.tf_.begin(),
	      bind(&a2time, _1, OmegaM, OmegaL, UnitTime_in_s));
    transform(nstar_particles.tf_.begin(),
	      nstar_particles.tf_.end(),
	      nstar_particles.tf_.begin(),
	      bind(&a2time, _1, OmegaM, OmegaL, UnitTime_in_s));
  }


  // The Springel files have parent IDs encoded in the ID's themselves.
  // We get the parent ID by erasing the higher bits
  transform(nstar_particles.id_.begin(),
	    nstar_particles.id_.end(),
	    back_inserter(nstar_particles.pid_),
	    _1 & 0x7FFFFFFF);

  transform(gas_particles.id_.begin(),
	    gas_particles.id_.end(),
	    back_inserter(gas_particles.pid_),
	    _1 & 0x7FFFFFFF);

  // do black hole particles
  read_dataset (*groups[bh], "BH_Mass", 
		bh_particles.mbh_, read_order[bh]);
  read_dataset (*groups[bh], "BH_Mdot", 
		bh_particles.mdotbh_, read_order[bh]);
  // bh particles have no parent id. SInce they can actually merge,
  // this doesn't make sense anyway.

  DEBUG(1,cout<< "Done reading snapshot" << endl;);

  return true;
}
			      
#endif
